[MiniMax-M3.1] MXFP4 index cache and MXFP4 KV cache (--index-cache-dtype fp4, --kv-cache-dtype fp4) - #2962

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

• ✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

• ✅ Run after the PR is approved and Pre Checkin passes
• ✅ Run immediately when an approval review is submitted
• ✅ Can be requested before approval with labels

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 2962 --add-label <label>

Round 2 pushed 13f413d. GSM8K: no regression within sampling variance. Performance follows in a separate comment.

Limit: the fixed score was measured on 2e700ec, not on 13f413d. The last commit only changes selector inputs the current aiter does not use (Triton fallback), so no change is expected, but it is not rerun end to end. GPU tests for padded graph replay, FP8-query top-k reuse and TBO plan keys are also not run yet.

Round 2 pushed 13f413d; performance below is for the first-repair tree, not 13f413d. AgentX c32 full 1200 s run, 0 errors. Final-head GPU tests and rerun are ongoing.

Measured tree b20dec0f4a24 (same tree as 2e700ec): 367.72 output tok/s/GPU (1470.88 / 4), 90th-percentile interactivity 108.04 tok/s/user (1000 / ITL p90 9.256 ms). aiter 5469555096; top-k selection falls back to Triton. #2770 reported 371.49 / 109.97 on a different machine, image and aiter, so it is a reference, not an A/B.

Round 3 pushed 13f413d; this is the final-head accuracy and GPU validation. GSM8K matches the earlier trees within sampling variance; GPU tests pass; padded graph replay ran on the real server. Final-head performance is still running and is not reported here.

Limits: TP2 / 2-index-head real replay is untested. The aiter pinned here has no small-k selector, so the new selector bounds ran through the Triton fallback and the aiter selector path is not exercised (its test is skipped). TBO was covered only by the unit test, not a real --enable-tbo server run. The padded-replay metadata tests are unit tests and do not themselves run a graph replay; the server trace below is the replay evidence.

Round 3 pushed 13f413d; final-head performance: 368.92 output tok/s/GPU, 106.78 tok/s/user at 90th-percentile ITL, 0 errors. Same config as the first-repair run (367.72 / 108.04); #2770's 371.49 / 109.97 is a different machine, image and aiter, so neither is a strict A/B.

ATOM_M3_INDEX_Q_FP8 is off (default) in this run, so FP8 index-query reuse is not measured. The FP8 main Q read by the PA kernel is a separate path and is no evidence for it.

Round 3 pushed 13f413d; reply to @valar-cop's review. Final-head GSM8K (accuracy) and AgentX c32 (performance) are posted. A real TP4 padded replay ran, but with the Triton selector: the pinned aiter has no small-k selector, so a padded replay with the aiter selector serving is not run. The three new inline threads are answered there; none is fixed in this head.

@ganyi1996ppo already reviewed 13f413d: #2962 (review) say @valar-cop re-review if you want a fresh one.

Round 3 pushed 13f413d; same-machine A/B for @valar-cop: output throughput -0.23%, p90 interactivity -0.83% against the pre-fix tree, within run-to-run spread. No code change.

@ganyi1996ppo already reviewed 13f413d: #2962 (review) say @valar-cop re-review if you want a fresh one.

@valar-cop re-review

ci triage for ececbeb:

ATOM SGLang Test passed. https://github.com/ROCm/ATOM/actions/runs/43346945790

ATOM vLLM Test passed. https://github.com/ROCm/ATOM/actions/runs/43346945786

ATOM Test passed. https://github.com/ROCm/ATOM/actions/runs/43346945757

Running Copilot Code Review passed. https://github.com/ROCm/ATOM/actions/runs/43237953246

Pre Checkin passed. https://github.com/ROCm/ATOM/actions/runs/43237930719

Round 3 pushed 13f413d; reply to @valar-cop's 05:01 review. A real TP2 / 2-index-head padded replay ran and GSM8K held. The selector-serving replay is not done: the small-k selector still falls back to Triton, so that run is pending. No fix is claimed here for the two new inline threads.

@ganyi1996ppo already reviewed 13f413d: #2962 (review) say @valar-cop re-review if you want a fresh one.

@valar-cop re-review

@valarLip re-read at da0c003; head is now 2f174c2 (fast-forward, 9 commits). Description and title are rewritten for this head. Nothing below was run on a GPU in this round (all cards were reserved); what is still open is said per item.

Your six blocking items (11-12)

1. aiter version probe: removed in 224ae84, the import is direct. The repo records no aiter minimum anywhere, so 2995ae4 / b60961e add a "Minimum aiter" note to docs/configuration_guide.md (#7014 the per-row selector, #7057 the fused writer, #7060 / #7061 the MXFP4 index scorers, #7056 MXFP4 pa_decode). Consequence, plainly: on an aiter without #7014 a bf16 / fp8 row of 131072+ tokens raises ImportError, where main fell back to Triton (Copilot flagged exactly this). I kept the direct import, as your rule says; the fallback would be one try/except, say if you want it.
2. getattr(config, "speculative_config", None): gone (224ae84).
3. One row rule: _metadata_rows serves prefill and decode (224ae84).
4. Boundaries: scorer allocation is in minimax_m3/index_buffers.py; _index_work_map_blocks, the index_mxfp4 bool and the kwargs helpers are gone (224ae84). Imports: nothing is imported per layer any more. The remaining function-local imports run at construction or startup, plus two cached accessors (flydsl_ops, _small_k_selector) that keep FlyDSL optional for bf16 / fp8 and index_buffers.py importable without aiter.
5. Title and description: rewritten (M3.1; FP4 index and FP4 KV both in scope; plain M3 with dense layers refused; the GPU validation says the last server-level run was at 55070fb).
6. Split: --spec-decode-kv-cache-dtype, its EngineArgs plumbing, the requested value, the "asked" half of resolve_draft_kv_dtype, the offload refusal on a differing draft dtype, its docs row and tests are in [Spec decode] --spec-decode-kv-cache-dtype: bf16/fp8 KV for a draft that has its own pool #3235 (based on main, independent of this PR, requested and settled values as two fields as you asked in 16aa94d update the link to the repo #3 / 97f7aac gpt_oss update: add fused_qk_rope_reshape_and_cache #9). They are gone from this PR (212f185). What stays here is what an fp4 target needs: a draft with a pool of its own is bf16 under --kv-cache-dtype fp4 (an fp4 draft has no layout), the family table that says whether a draft has its own pool, V4's fp8 -> bf16 fallback ordering in Config, and the one settled field draft_kv_dtype. resolve_draft_kv_dtype and the EAGLE / DSpark builder changes are the same code in both PRs, so they can leave this one only if [Spec decode] --spec-decode-kv-cache-dtype: bf16/fp8 KV for a draft that has its own pool #3235 lands first. Order: the two are independent; whichever lands second is rebased. This PR is 49 files, +4198 / -402 (was +4284 / -402): the split alone does not make it small.

Also changed

• B12: c48f677 (thread answered). B15, the mapping half: 70bd1ad.
• 11-10 update readme for deepseek #6: an omitted --index-cache-dtype beside --kv-cache-dtype fp4 now defaults to fp4; an explicit bf16 / fp8 is still refused (2f174c2, test test_an_fp4_kv_without_an_index_flag_gets_the_fp4_index). This reverses 4e4dca9.
• B13: the allow-list is still in config.py; its message and docstring now say MiniMax-M3.1 (2f174c2). Moving it needs a per-model hook for the generic refusal: not done.
• Three small commits (75ffc24, 0c514f9, 2e3790a): two duplicate tests, a one-line wrapper, a long docstring removed.

Still open, not changed

• B1, B3, B4: the scorer works at the model-length width on every step (eager, graph, TBO ubatch); eager decode plans still accumulate per row count. Needs a measurement and, for B3, an aiter change.
• B5: plan_segment and the segment loop stay; they need MXFP4 scale support in run_pa_decode.
• B7, B8, the KvLayout half: the pool still takes kv_dtype and a layout; about 80 lines and 20 test sites, and main's transfer_pages reads pool.kv_dtype.
• The two skip-index writers; the DP test (one rank prefilling while another decodes); T9, the score workspace (512 MiB at TP4 / 1M context, 4x at TP1).
• CI: the MiniMax-M3-MXFP4 accuracy job runs --kv_cache_dtype fp8, so a green CI says nothing about the fp4 paths. An fp4 job needs a GPU-runner decision from the maintainers; I added none.
• GPU validation of this head: the last server-level run was 55070fb (with the mxfp4 spelling of the flag); since then only kernel-level checks on one card. The rerun commands are ready.
• Seven threads (B1, B3, B4, B5, B7, B8, B13) show as resolved because we resolved them with "tracked for follow-up". That was premature: they are open work, and I answered each below; I did not reopen them.

Corrections to things I said earlier (the threads are not edited)

• 6647558639 (11-05): "an eager main decode step scores at the batch's block count" is false; every step scores at cdiv(max_model_len, 128) (the shrink was removed in 7bbdd07). And "tests/test_minimax_m3_skip_index_writer.py is written" is false: there is no such file.
• T28 (B8): "MhaKvPool no longer stores kv_dtype" is false at this head: the merge of main restored self.kv_dtype, which main's transfer_pages reads.
• 5771103020 item 10: "tests/test_v41_geometry_parity.py reads that spelling via AST": there is no such file and no test reads it. model_type == "deepseek_v41_text" is main's own spelling in five places; the PR adds one. Nothing changed.
• T26 (B6): the reply names decode_rows_batch, inlined later (97f7aac).
• 6731087765 item 2: "same rule for every index cache" was superseded; the rule is MXFP4-only and fp4 KV with another index is refused.
• T22 (B2), T30 (B10), T31 (B11), T39, T46 / T52: the threads say "follow-up / not changed / still open", but the code was changed later (620692b, fd64ce3, 785e40f, 68fc39c, 4cc42e7).
• Reply 6786465138 item 3: Config also imports aiter for V4's fp8 -> bf16 fallback (get_gfx), not only for M3.

@valar-cop I did not ask you to re-review after 224ae84, 5c76e58, da0c003 and this push; this is that request.

@valarLip re-read of your 11-12 11:22 review (5873510797) at 2808171. New head: 015f202, pushed once. It carries the changes below; the unpushed test fix is item 1.

1. fp4 paths on a GPU at this head
Run on 4x MI355X (gfx950), the code of 015f202 (plus one test-only commit, below), aiter from the rocm/atom-dev:nightly_202611041513 image (commit 627461276).

• gfx950 pytest files (test_minimax_m3_emit_pad_row, _index_topk, _mxfp4_index, _mxfp4_write, _mxfp4_attention, _mxfp4_decode_batch): 187 passed, 32 skipped, 1 failed. The 32 skips are the gfx1250 launch-policy regression parametrisations in test_minimax_m3_index_topk.py, as intended on gfx950. The failure is a bug in the test, not in the product: _run in test_minimax_m3_mxfp4_index.py allocated the scorer buffers inside a CUDA-graph capture (hipErrorStreamCaptureUnsupported) in the graph variant of the padded-step test. Fixed in 7fe048d (test only, +3/-2); that file then gives 10 passed. That commit is not pushed yet (I pushed once this round); it goes with the next push.
• GSM8K, M3.1 TP4, --kv_cache_dtype fp4 --index-cache-dtype fp4, DSpark with 7 drafts, --block-size 128, 5-shot chat, temperature 0, first 200 problems only: flexible 0.960 (stderr 0.014), strict 0.900 (0.021). Startup log: 60 h1d128fp4 rows in the target pool, the DSpark draft pool allocated separately. Cumulative [MTP Stats] after the run: 4.52 tokens per forward, acceptance 50.35%. It is 200 of 1319 questions, and the server was launched without the recipe's --default-chat-template-kwargs thinking flag, so it is not like-for-like with the recipe's 0.955 / 0.914 (fp8 KV, full set).
• --enable-tbo all smoke, fp4 KV + fp4 index, without DSpark. With DSpark the server refuses to start: NotImplementedError: MiniMax M3 sparse TBO decode with max_q_len > 1 is not implemented (aiter_attention.py, the same check as on main). Without it the server starts (TBO enabled: model wrapped with UBatchWrapper), captures graphs and serves. GSM8K, first 100 problems: flexible 0.93 / strict 0.88. The same checkout without TBO, two runs: 0.96 / 0.93 and 0.95 / 0.92. On strict outcomes the TBO run differs from the two non-TBO runs in 11 and 6 of 100 problems, the two non-TBO runs from each other in 9; raw text is identical in only 13 to 16 of 100 in every pair. So the TBO result is inside the run-to-run spread at N=100; that is a smoke, not a parity proof. Split-bucket graph replay under DP is not covered.
• Copilot's EAGLE3 draft-dtype test (test_an_eagle3_draft_attention_is_built_at_the_dtype_of_its_pool): passes. It stubs EagleProposer, imports aiter, and allocates no pool and runs no forward, so it needs no GPU; I ran it on a CPU host.
  Observation: every fp4 server start here printed KV cache allocation mismatch: expected=169GB, actual=136GB (16-20%), with and without DSpark and TBO; I did not run bf16 or fp8 for comparison, so I do not know whether it is specific to fp4.
  Not run: the other cache combinations (fp8 KV + fp4 index, bf16 index), a DP step with one rank prefilling while another decodes, a request of 131072 tokens or more.

2. Land #3235 first, then rebase
This branch is already the flag-less form of #3235: no --spec-decode-kv-cache-dtype, no "asked" field, and resolve_draft_kv_dtype(target_dtype, *, own_pool) is the one-line fp4 -> bf16 rule. What is still the same code is the mechanism (draft_owns_kv_pool, DRAFT_POOL_FAMILIES, SpeculativeConfig.draft_kv_dtype and the proposer, calibration and hash edits). I cannot rebase without a force push, so my plan is: when #3235 merges, merge main into this branch (a merge commit), take main's side in the draft files and re-add the fp4 rule on top of main's resolve_draft_kv_dtype. I did that as a dry run against the #3235 head a3802ba: it conflicts in 8 files, and what this PR keeps after the merge in the draft area is the fp4 clause in the resolver, one clause in the offload refusal, a few test rows and one doc sentence. The CPU tests of both areas on the merge result show only failures that also occur on this head without the merge (listed below). Until #3235 merges both PRs carry the mechanism; I did not stack the flag into this PR, which would add about 460 lines to its diff.
Asking: is merging main into this branch after #3235 lands acceptable, or do you want it stacked now?

3. Design items
Fixed in 015f202:

• KvGeometry.layout is typed KvLayout | None (pool_rows.py).
• The shared MHA builder no longer imports Mxfp4KvLayout or tests the dtype in _module_kinds: SparseMHAPagedAttentionImpl.kv_layout declares the layout (None unless fp4) and the builder reads module.impl.kv_layout. The builder's _index_cache_spec still names the MXFP4 index page (PAGE_BYTES, Mxfp4KvLayout().check); I did not change that.

Not changed, asking for your explicit OK to defer each, with an issue I would open on your OK (nothing is created yet):

• B5, segment loop (sparse_attn_mxfp4.py:108, about 28 lines) and plan_segment (base_attention.py:502, about 7 lines): replacing them by a dtype branch in run_pa_decode is about -35 lines here, but it needs MXFP4 K/V scale arguments in run_pa_decode (a change in base_attention.py and in aiter's pa_decode entry) and a gfx950 re-run of the sparse PA tests and GSM8K, since it is the decode path of every fp4 request. Proposed issue: "MiniMax-M3.1 MXFP4 KV: route sparse PA decode through run_pa_decode and drop the segment plan loop".
• B7/B8, MhaKvPool takes kv_dtype and layout (mha_kv_pool.py:204-208) and transfer_pages reports self.kv_dtype (:379): dropping kv_dtype is about 35 source lines and 20 test call sites, and transfer_pages must then take the dim and dtype of each region from the layout; for an fp4 pool that changes what is reported (today head_dim and uint8), which a P/D pair has to agree on. Proposed issue: "MhaKvPool: take only the KV layout; the layout reports the transfer unit".
• B13, per-model fp4 allow-list _check_fp4_kv (config.py:1948): "every other model refuses fp4" has to live in shared code; moving the M3 and V4.1 knowledge out needs a new per-model hook (about 40 lines). Proposed issue: "A per-model supports_fp4_kv hook instead of _check_fp4_kv in config.py".
• The two skip-index writers (the Triton writer for non-fp4 skip layers, and aiter's skip_index_branch for fp4 KV in attention_mha.py): one writer needs a gfx950 comparison of the two on the same input and a decision on which page format a skip layer under fp8 KV + fp4 index uses. Proposed issue: "M3 skip-index layers: one fused writer for every KV dtype".

4. Docs
docs/configuration_guide.md now lists, in the kv_cache_dtype row, decode context parallelism (--decode-context-parallel-size > 1) and a speculative draft that would keep FP4 KV (an MTP draft shares the target's pool; a DSpark or EAGLE3 draft with its own pool is bf16 and runs), and in the index_cache_dtype row the DCP refusal (015f202). The rules are the ones in models/minimax_m3/config.py:69 and _mxfp4_kv_refusal.

CPU checks at 015f202: black and ruff clean on the changed files. Two test files cannot be collected on that host (tests/test_dspark.py, test_dspark_swa_fp8_2buff.py: aiter stub). The same failures occur without this commit's changes: test_spec_lossless_staging.py, test_paged_attention_dispatch.py::TestDenseDecodeSplits, test_deepseek_v4_wo_a_dequant.py, test_deepseek_v41_mono_hi_gate.py, test_minimax_m3_flydsl_optional.py (on the CPU host the image's aiter has no aiter.ops.triton.fused_kv_cache in a bare interpreter; it imported fine on the GPU host). Size: 47 files, +2837/-412 before, +2848/-420 now (the new lines are the doc sentences and the impl's kv_layout; the builder lost its M3 branch).

@valarLip follow-up to my comment above. New head: bd8fd40.

• My previous push (015f202) broke tests/test_layout_packages.py: pool_layout/pool_rows.py imported atom for the KvLayout annotation, which that package promises not to do. This is our regression, found after I posted; the "Run non-GPU unit tests" job failed because of it.
• Fix (bd8fd40): KvGeometry.layout is annotated with a one-attribute Protocol (tag) defined next to KvGeometry, so there is no import. tests/test_layout_packages.py is unchanged and passes (31 passed; it fails on 015f202).
• 7fe048d is now included, so the unpushed note in my earlier comment no longer applies: the graph variant of the padded-step index test no longer allocates inside the capture (test_minimax_m3_mxfp4_index.py, 10 passed on gfx950).
• bd8fd40 also drops the three comment corrections of 015f202 (deployment.py, eagle_proposer.py, draft_kv.py; they belong to [Spec decode] --spec-decode-kv-cache-dtype: bf16/fp8 KV for a draft that has its own pool #3235) and shortens the docs wording. Against main the PR is now +2841/-413 in 47 files (it was +2837/-412 before this round).
• Checks on a CPU host: black and ruff clean on the changed files; the touched-area tests give the same failures as before this round (environment ones, also present at 2808171).

@valarLip re-read of your review 5875809244 at bd8fd40. New head: ececbeb, pushed once.

1. The getattr in the shared builder: fixed in ececbeb. PagedAttentionImpl.kv_layout = None is declared once on the impl base, the builder reads module.impl.kv_layout, and the _PagedAttention stub in tests/test_draft_kv_builder.py has an impl.

2. Order with #3235: #2962 lands first and #3235 after it, as agreed with the maintainers; nothing to merge here for #3235.

3. Follow-up issues: opened, and linked in the PR description under "Follow-ups" (the rest of the description is unchanged).

• B5, segment loop and plan_segment: MiniMax-M3.1 MXFP4 KV: route sparse PA decode through run_pa_decode and drop the segment plan loop #3260
• B7/B8, MhaKvPool kv_dtype and layout: MhaKvPool: take only the KV layout (drop kv_dtype); the layout reports the transfer unit #3261
• B13, _check_fp4_kv: Per-model fp4 KV support hook instead of _check_fp4_kv in shared config.py #3262
• the two skip-index writers: MiniMax-M3 skip-index layers: one fused writer for every KV dtype #3263

4. KV cache allocation mismatch: not fixed, and I have not been able to run the fp8 comparison: all eight GPUs are in use by someone else's TP8 server, and I am not allowed to take them. I will run the same server with --kv-cache-dtype fp8 when four cards are free. What the fp4 logs already show (DSpark and no-DSpark runs, TP4):

• The check compares the rank's own plan (pool_plan.total_reserved_bytes) with the bytes the rank allocated. The ranks sized different budgets: non_torch was 45.8 / 45.5 GB on two ranks and 17.2 / 16.9 GB on the other two, giving 98454 / 98647 / 118062 / 118255 blocks. Every rank then allocated 98454 blocks, the count the engine took from get_num_blocks (engine_core.py:146, then allocate_kv_cache(num_blocks)), which is the smallest of the four here.
• The allocated bytes are exactly that count times block_bytes (143.656 GiB = 98454 x 1566720 B), and only the two ranks whose own plan was larger (172.3 / 172.5 GiB) print the warning. So in these runs fp4 is not losing blocks it reserved: the pool is priced the same on the plan and the allocation side. The open question is why non_torch differs by about 28 GB between ranks; I have not found it, and I do not know yet whether fp8 shows the same.

5. Test leftovers: fixed in ececbeb. test_what_still_passes_and_v41_keeps_its_fp4_kv now has pytest.raises for _check_fp4_kv("fp4", is_m3=False, is_v41=False); the docstring of test_minimax_m3_emit_pad_row.py no longer says GPU-PENDING (the gfx950 variant ran on MI355X and passed).

Copilot's two unpublished findings

• (a) Index columns in the FP4 KV / Ulysses writer path (review 5874692436): real, in one configuration. The aiter fused writer checks qkv.size(1) == (nq + 2*nkv + niq + 1) * 128 and requires num_index_heads > 0 even with skip_index_branch=True (csrc/kernels/fused_qknorm_idxrqknorm.cu, expected_row), while SparseMHAPagedAttentionImpl.qkv_fields drops the index columns of a top-k reusing layer, so under --sequence-parallel-size > 1 the first reusing layer would fail that check. Fixed in ececbeb by refusing --kv-cache-dtype fp4 with sequence parallelism at startup (4 lines, a test row and a doc sentence). Not run on a GPU; the check itself is what I read in aiter, not something I executed.
• (b) Split prefill passing validation and failing at runtime (review 5875178617): not confirmed. I looked at the prefill metadata of a split (TBO) prefill: make_sparse_prefill_metadata drops the work map and keeps only the score workspace, which both micro-batches use one after the other on the compute stream, and n_valid_column_per_row is allocated per call. I found no path that fails, but I did not run --enable-tbo (prefill only) with fp4 on a GPU, so I cannot call it not real; it is pending a GPU, and the finding has no details to check against.

CPU checks at ececbeb: black and ruff clean on the changed files; tests/test_layout_packages.py and the touched-area tests give the same failure set as before (environment ones, also present before this round). Size against main: 47 files, +2851/-413 (+2841/-413 at bd8fd40): +10 lines, of which the sequence-parallelism refusal with its test and doc sentence is 6 and the two test leftovers 4.

@valarLip answer to your item 4 (the KV cache allocation mismatch), measured. No code change, nothing pushed for it.

Cause: not the fp4 pool, and not a bug in this PR. My earlier GPU runs set HIP_VISIBLE_DEVICES assuming HIP index order equals PCI order. It does not on that machine (HIP index follows the KFD node order), so two of the four ranks ran on cards that another job on the same machine was also using, with about 28 GB of its own memory resident there. That memory shows up in non_torch ((total - free) - memory_reserved, model_runner.py:2051) on exactly those two ranks, 45.5 / 45.8 GB against 17.2 GB on the other two, and the difference is the 28 GB. The two ranks with the larger budget then print the warning, because every rank allocates the block count the engine took from rank 0 (async_proc.py:613, engine_core.py:146), not its own.

Re-run on four free cards (the same M3.1 TP4 + DSpark server as the GSM8K run, ececbeb, only the KV and index dtype changed; HIP indices derived from the KFD order, the other cards idle):

Post-init device memory is 251.9 to 252.6 GB on every rank in both runs. So fp4 prices the pool the same way on the plan and the allocation side, and loses no blocks it reserved; fp4 holds 1.67x the blocks of fp8 for the same memory.

Two corrections to my previous comment: the 98454 blocks that every rank allocated was rank 0's count, not "the smallest of the four" (it was the smallest only because rank 0 sat on a shared card), and the GSM8K / TBO numbers at 015f202 above came from that same shared-card setup. Accuracy does not depend on it; I did not measure throughput.

Not changed: the engine uses rank 0's count without taking a minimum across the TP ranks. With uniform ranks the counts differ by 0.2% and the 2% safety margin covers it; I left it alone.

@valarLip re-read of your review 5884830762. New head: bc61ed4 (a merge commit: main merged into this branch, no rebase, no force push; the branch was conflicting with main after #2581).

Merge conflicts (5 files, all from the Kimi-K3 DSpark GQA draft, #2581, which touches the same draft-pool code):

• atom/spec_decode/draft_kv.py: kept this branch's family table (draft_owns_kv_pool, DRAFT_POOL_FAMILIES) and the draft dtype settled by Config, took main's DraftKvPoolSpec, draft_kv_dtype() and the pool-spec kwargs. A draft that declares a pool spec on a family without a draft pool is refused, as on main, now by the table instead of DRAFT_OWNS_KV_POOL. The scale binding follows the pool (pool.scale_fields).
• atom/model_ops/attentions/mha_kv_pool.py: kept the KvLayout, took main's flash_layout (checked and priced as SHUFFLE unless FLASH: no scale planes, fa_views).
• atom/model_ops/attentions/backends.py: docstring only.
• tests/test_draft_kv_builder.py, tests/test_dspark.py: both sides kept; the stubs follow the table and the settled dtype (tests/test_kimi_k3_dspark_gqa.py needed the same).

1. Split-prefill finding: refused, not run. I did not reproduce it: my earlier TBO runs used short prompts and I did not verify that a prefill split happened, and no GPU run was done for this round. The assert in _metadata_rows compares the sparse metadata's row count with q.shape[0], and under prefill context parallelism a rank's rows are a shard. So validate_runtime_config now refuses --prefill-context-parallel-size > 1 with either MXFP4 cache (test rows for both, doc line in configuration_guide.md), in bc61ed4.

3. Rename: test_what_still_passes_and_v41_keeps_its_fp4_kv is now test_fp4_kv_is_m3_and_v41_only (no other references).

2. CI: the offline smoke failure was the runner's docker: unauthorized, not this PR. I will rerun only that job once if it fails again after this push; the heavy accuracy jobs are still pending.

CPU checks: black clean; ruff clean on the files this PR changed; tests/test_layout_packages.py and the touched-area tests give the same failures as before the merge (environment ones), plus test_kimi_k3_dspark_gqa.py::test_the_builder_declares_draft_slots_only_under_dcp[8-True], which needs a GPU for a pinned buffer.